package com.kklsqm.langchain4j.demo;

import org.springframework.http.*;
import org.springframework.web.client.RestTemplate;

import java.util.*;


public class DashScopePortraitDemo {

    public static void main(String[] args) throws InterruptedException {
        String apiKey = "sk-7fc88c005edb46928c4ef180268169a2";
        String imageUrl = "https://img.bosszhipin.com/beijin/upload/avatar/20240427/607f1f3d68754fd0ba6c682a237670d7542f37142655d384b611e718b4b9b78a7b71959ff6ab8847_s.png";
        String templateUrl = "https://p3-pc-sign.douyinpic.com/tos-cn-i-0813c001/ooBgUDgPlFYahWAEqAiijAGZ5lI1PDsA2AIEm~tplv-dy-aweme-images:q75.webp?biz_tag=aweme_images&from=327834062&lk3s=138a59ce&s=PackSourceEnum_SEARCH&sc=image&se=false&x-expires=1749279600&x-signature=uvyuSGC6P1bh7cATmRNXHPSqQwI%3D";

        String taskId = generatePortrait(apiKey, imageUrl, templateUrl);
        System.out.println("任务提交成功，ID：" + taskId);

        String resultUrl = queryPortraitResult(apiKey, taskId);
        System.out.println("写真生成成功，图片地址：" + resultUrl);
    }


    public static String generatePortrait(String apiKey, String imageUrl, String templateUrl) {
        String url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/album/gen_potrait";

        RestTemplate restTemplate = new RestTemplate();

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(apiKey);
        headers.set("X-DashScope-Async", "enable");

        Map<String, Object> requestBody = new HashMap<>();
        requestBody.put("model", "facechain-generation");

        Map<String, Object> parameters = new HashMap<>();
        parameters.put("style", "train_free_portrait_url_template");
        parameters.put("n", 1);
        parameters.put("skin_retouch", true);
        requestBody.put("parameters", parameters);

        Map<String, Object> input = new HashMap<>();
        input.put("user_urls", List.of(imageUrl)); // 头像
        input.put("template_url", templateUrl);    // 自定义风格模板
        requestBody.put("input", input);

        HttpEntity<Map<String, Object>> requestEntity = new HttpEntity<>(requestBody, headers);

        ResponseEntity<Map> response = restTemplate.postForEntity(url, requestEntity, Map.class);

        if (response.getStatusCode() == HttpStatus.OK) {
            Map<String, Object> output = (Map<String, Object>) response.getBody().get("output");
            return output.get("task_id").toString();
        } else {
            throw new RuntimeException("调用失败，状态码: " + response.getStatusCode());
        }
    }


    public static String queryPortraitResult(String apiKey, String taskId) throws InterruptedException {
        String queryUrl = "https://dashscope.aliyuncs.com/api/v1/tasks/" + taskId;

        RestTemplate restTemplate = new RestTemplate();
        HttpHeaders headers = new HttpHeaders();
        headers.setBearerAuth("Bearer "+apiKey);

        HttpEntity<Void> entity = new HttpEntity<>(headers);

        int retry = 0;
        while (retry < 20) { // 最多轮询20次
            ResponseEntity<Map> response = restTemplate.exchange(queryUrl, HttpMethod.GET, entity, Map.class);
            if (response.getStatusCode() == HttpStatus.OK) {
                Map<String, Object> body = response.getBody();
                String status = (String) ((Map<Object, Object>)body.get("output")).get("task_status");
                if ("SUCCEEDED".equals(status)) {
                    Map<String, Object> output = (Map<String, Object>) body.get("output");

                    // 获取results列表（每个元素是一个Map）
                    List<Map<String, Object>> results = (List<Map<String, Object>>) output.get("results");

                    if (results != null && !results.isEmpty()) {
                        // 获取第一个结果的url字段
                        return (String) results.get(0).get("url");
                    } else {
                        throw new RuntimeException("生成成功但未返回有效结果");
                    }
                } else if ("FAILED".equals(status)) {
                    System.out.println("任务执行失败，失败原因：" + ((Map<Object, Object>)body.get("output")).get("code"));
                    System.out.println("任务执行失败，失败原因：" + ((Map<Object, Object>)body.get("output")).get("message"));
                    throw new RuntimeException("任务执行失败");
                } else {
                    System.out.println("当前状态: " + status + "，等待3秒继续...");
                    Thread.sleep(3000);
                    retry++;
                }
            } else {
                throw new RuntimeException("查询任务失败，状态码：" + response.getStatusCode());
            }
        }

        throw new RuntimeException("任务超时未完成");
    }

}
